import torch
import numpy as np
from PET.data_handle.template import HardTemplate
from PET.pet_config import ProjectConfig
from transformers import AutoTokenizer
from datasets import load_dataset
from functools import partial


def convert_example(examples: dict, tokenizer, max_seq_len: int, max_label_len: int,
                    hard_template: HardTemplate, train_mode=True, return_tensor=False) -> dict:
    """
    将样本数据转换为模型接收的输入数据。
    Args:
        examples (dict): 训练数据样本, e.g. -> {"text": ['手机	这个手机也太卡了。','体育	世界杯为何迟迟不见宣传',...]}
        max_seq_len (int): 句子的最大长度，若没有达到最大长度，则padding为最大长度
        max_label_len (int): 最大label长度，若没有达到最大长度，则padding为最大长度
        hard_template (HardTemplate): 模板类。
        train_mode (bool): 训练阶段 or 推理阶段。
        return_tensor (bool): 是否返回tensor类型，如不是，则返回numpy类型。
    Returns:
        dict (str: np.array) -> tokenized_output =
        {
            'input_ids': [[1, 47, 10, 7, 304, 3, 3, 3, 3, 47, 27, 247, 98, 105, 512, 777, 15, 12043, 2], ...],
            'token_type_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], ...],
            'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], ...],
            'mask_positions': [[5, 6, 7, 8], ...],
            'mask_labels': [[2372, 3442, 0, 0], [2643, 4434, 2334, 0], ...]
            }
    """
    tokenized_output = {
        'input_ids': [],
        'token_type_ids': [],
        'attention_mask': [],
        'mask_positions': [],
        'mask_labels': []
    }

    for i, example in enumerate(examples['text']):
        if train_mode:
            label, content = example.strip().split('\t')
        else:
            content = example.strip()

        inputs_dict = {
            'textA': content,
            'MASK': '[MASK]'
        }
        encoded_inputs = hard_template(
            inputs_dict=inputs_dict,
            tokenizer=tokenizer,
            mask_length=max_label_len,
            max_seq_len=max_seq_len
        )
        tokenized_output['input_ids'].append(encoded_inputs['input_ids'])
        tokenized_output['token_type_ids'].append(encoded_inputs['token_type_ids'])
        tokenized_output['attention_mask'].append(encoded_inputs['attention_mask'])
        tokenized_output['mask_positions'].append(encoded_inputs['mask_position'])

        if train_mode:
            label_encoded = tokenizer(text=label, add_special_tokens=False)
            # print('label_encoded:', label_encoded)
            label_encoded = label_encoded['input_ids'][:max_label_len]  # 裁剪
            # print('label_encoded:', label_encoded)
            label_encoded = label_encoded + [tokenizer.pad_token_id] * (max_label_len - len(label_encoded))  # 补齐
            # print('label_encoded:', label_encoded)
            tokenized_output['mask_labels'].append(label_encoded)

    for key, value in tokenized_output.items():
        if return_tensor:
            tokenized_output[key] = torch.LongTensor(value)
        else:
            tokenized_output[key] = np.array(value)

    return tokenized_output


if __name__ == '__main__':
    pc = ProjectConfig()
    pc.pre_model = '../../pre_model/bert-base-chinese'
    pc.train_path = '../data/train.txt'
    pc.dev_path = '../data/dev.txt'
    pc.prompt_file = '../data/prompt.txt'
    pc.verbalizer = '../data/verbalizer.txt'

    tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)
    hard_template = HardTemplate(prompt='这是一条{MASK}评论：{textA}')

    # train_dataset = load_dataset('text', data_files=pc.train_path)
    # print(train_dataset)
    # print(train_dataset['train'])
    train_dataset = load_dataset('text', data_files=pc.train_path, split='train')
    # print(train_dataset)
    # # print(train_dataset['text'])
    # examples = train_dataset

    # # examples = {"text": ['手机	这个手机也太卡了。', '体育	世界杯为何迟迟不见宣传']}
    # # print(examples)
    # tokenized_output = convert_example(examples=examples,
    #                                    tokenizer=tokenizer,
    #                                    max_seq_len=30,
    #                                    max_label_len=2,
    #                                    hard_template=hard_template,
    #                                    train_mode=True,
    #                                    return_tensor=True)
    # print('tokenized_output:', tokenized_output)

    convert_fun = partial(convert_example, tokenizer=tokenizer,
                          max_seq_len=30, max_label_len=2, hard_template=hard_template)
    # tokenized_output = convert_fun(examples)

    dataset = train_dataset.map(convert_fun, batched=True)
    print('dataset:', dataset)
    for value in dataset:
        print(value)
        print(value.keys())
        print("value['text']:", value['text'])
        print("value['input_ids']:", value['input_ids'])
        print("value['token_type_ids']:", value['token_type_ids'])
        print("value['attention_mask']:", value['attention_mask'])
        print("value['mask_positions']:", value['mask_positions'])
        print("value['mask_labels']:", value['mask_labels'])
        break

        # value['text']: 电脑	(1)这款笔记本外观感觉挺漂亮的，分量吗，对我来说不算沉。 (2)安装了WindowsXP系统后，运行的速度挺快。发热量没有想象中那么大。可能尚未运行很耗资源的程序，没有感到内存的弊病。不过，1G的内存确实有点小。 (3)附赠的包很不错，挺有手感的。但是附赠的鼠标实在是太小了，幸好同时订了一个双飞燕的鼠标哟。
        # value['input_ids']: [101, 6821, 3221, 671, 3340, 103, 103, 6397, 6389, 8038, 113, 122, 114, 6821, 3621, 5011, 6381, 3315, 1912, 6225, 2697, 6230, 2923, 4023, 778, 4638, 8024, 1146, 7030, 102]
        # value['token_type_ids']: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
        # value['attention_mask']: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
        # value['mask_positions']: [5, 6]
        # value['mask_labels']: [4510, 5554]

        # input_ids：编码后的输入 ID。
        # token_type_ids：句子类型 ID，是否是一句话。
        # attention_mask：注意力掩码。
        # mask_positions：[MASK] 的位置。
        # mask_labels：标签的编码（仅在训练模式下）。
